from qutip import *
import numpy as np
import matplotlib.pyplot as plt
sys.path.append('D:\Dropbox\Dokumente\PI3\Code\Python\quantum-tools')
import qutip_enhanced
qte = qutip_enhanced.QutipEnhanced()

b = Bloch3d()

def alpha(n, nu_rf, tau, length_mus_pi):
    return int(np.ceil(n/2.))*(2*np.pi*nu_rf*(tau + length_mus_pi) + np.pi)

def h(nu_rf, hf, omega_e=0, phase_e=0, omega_n=0, phase_n=0):
    a = 2*np.pi*nu_rf/2.
    b = 2*np.pi*(nu_rf - hf)/2.
    c = 0.5*2*np.pi*omega_e*np.exp(1j*phase_e)
    d = 0.5*2*np.pi*omega_n*np.exp(1j*phase_n)
    dcn=0
    h = Qobj([[a + dcn, 0,          c.conj(), 0],
             [0,        -a + dcn,   0,        c.conj()],
             [c,        0,          b+dcn,    d.conj()],
             [0,        c,          d,        -b+dcn]])
    h.dims = [[2, 2], [2, 2]]
    return h

def rho0():
    rho0_e = ket2dm(basis(2, 0))
    rho0_n = ket2dm(basis(2, 0))
    return tensor(rho0_e, rho0_n)

def evolve(rho, length_mus, **kwargs):
    options = Odeoptions(nsteps=100000)
    if length_mus != 0.0:
        rho = mesolve(h(**kwargs), rho, np.linspace(0, length_mus, 2), [], [], options=options).states[-1]
    # b.add_states(rho.ptrace(1))
    return rho

def KDD5(rho, shift, tau, omega_e, omega_n, nkdd5=0, length_mus_pi=0, **kwargs):

    #tau
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+0, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)

    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=np.pi / 6. + shift, **kwargs)

    #two tau
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+1, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+2, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+3, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+4, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=np.pi / 2. + shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+5, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+6, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+7, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+8, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=np.pi / 6. + shift, **kwargs)
    #tau
    rho = evolve(rho, tau, omega_n=omega_n, phase_n=alpha(10*nkdd5+9, tau=tau, nu_rf=kwargs['nu_rf'], length_mus_pi=length_mus_pi), **kwargs)
    return rho


def sequence(**kwargs):
    rho = rho0()

    rho = evolve(rho, length_mus=.25/kwargs['omega_e'], omega_e=kwargs['omega_e'], hf=kwargs['hf'], nu_rf=kwargs['nu_rf'])

    rho = KDD5(rho, shift=0.0, nkdd5=0, **kwargs)
    rho = KDD5(rho, shift=np.pi / 2., nkdd5=1, **kwargs)
    rho = KDD5(rho, shift=0.0, nkdd5=2, **kwargs)
    rho = KDD5(rho, shift=np.pi / 2., nkdd5=3, **kwargs)

    rho = evolve(rho, length_mus=.25/kwargs['omega_e'], omega_e=kwargs['omega_e'], hf=kwargs['hf'], nu_rf=kwargs['nu_rf'])
    return rho

def tau_plus_pi_dur_spectrum():
    _tau_ = 60.
    _omega_e_ = 1/0.2
    _omega_n_ = 1/(40*_tau_)
    _length_mus_pi_ = 0.25/_omega_e_
    _hf_ = 0.089
    _nu_rf_ = _hf_
    kwargs = dict(tau=_tau_, omega_e=_omega_e_, omega_n=_omega_n_, hf=_hf_, length_mus_pi=_length_mus_pi_, nu_rf=_nu_rf_)
    b.add_states(sequence(**kwargs).ptrace(0))
    b.show()


    # b.add_states()


    # for _nu_rf_ in _nu_rf_list:
    #     kwargs['nu_rf'] = _nu_rf_
    #     exp_z.append(-expect(jmat(0.5, 'z'), sequence(**kwargs).ptrace(1)))
    # plt.figure()
    # plt.plot(_nu_rf_list, exp_z)
    # plt.show()

tau_plus_pi_dur_spectrum()

# exp_z = []
# for _tau_ in tau_list:
#     kwargs = dict(tau=_tau_, omega_e=_omega_e_, omega_n=_omega_n_, hf=_hf_, nu_rf=_nu_rf_)
#     exp_z.append(-expect(jmat(0.5, 'z'), sequence(**kwargs).ptrace(0)))
#
# from scipy.optimize import curve_fit
#
# def exponential(t, a,b,c):
#     return a*exp(-b*t) + c
# import matplotlib.pyplot as plt
# plt.plot(tau_list, exp_z)
# popt, pcov = curve_fit(exponential, tau_list, exp_z, p0=[0.5, 4e-3, 0])
# print 40/popt[1]
# plt.plot(tau_list, exponential(np.array(tau_list), *popt))
# plt.show()
# res_arr = np.loadtxt("res_arr{}.dat".format(_hf_))
#
# print "omega_n: {}".format(_omega_n_)
# plt.figure()
# for i in [10, 13,15,18,19]:
#     plt.plot(nu_rf_list, res_arr[:, i], label="{}".format(omega_n_list[i]))
#     plt.legend()
# plt.show()